/**
 * Interactive heatmap using D3
 * 
 * Matthew Iyer
 * Nov 20, 2013
 */

function Heatmap(config) {
	// dimensions
    var width = config.width || 960;
    var height = config.height || 960;
    var padding = parseInt(config.padding) || 5;
    var margin = config.margin || { top: 25, right: 300, bottom: 300, left: 25 };
	// heatmap data source
    var rowLabelField = config.rowLabelField || 'name';
    var colLabelField = config.colLabelField || 'name';
	var heatmapValueField = config.heatmapValueField || 'global_nes';
	var rowLabelFunc = config.rowLabelFunc;
    var url = config.url || "{{ url_for('json_request_heatmap_data') }}";
    // heatmap colors
    var colorDefault = config.colorDefault || "#888888";
    var defaultColorScale = d3.scale.linear()
    	.domain([-20, 0, 20])
    	.range(["blue", "white", "red"])
    	.clamp(true);
    var colorScale = config.colorScale || defaultColorScale;
    // animation
    var duration = 200;
    // tool tips
    var toolTipHtmlFunc = config.toolTipHtmlFunc || function(d) { return d.value.value; };
    
    // private data
    var numControlButtons = 2;
    var innerWidth = width - padding - margin.left - margin.right;
    var innerHeight = height - padding - margin.top - margin.bottom;
    var rowIds = []; // ordered row ids
    var colIds = []; // ordered column ids
    var rowMap = d3.map(); // map row id to row data
    var colMap = d3.map(); // map col id to col data
    var cellMap = d3.map(); // map 'row,col' key to cell data
    var toggle = true;
    
    // private d3 variables
    var rootSelection = null;
    var gridScaleRow = d3.scale.ordinal();
    var gridScaleCol = d3.scale.ordinal();
    // column and row dragging behavior
    var drag = d3.behavior.drag()
		.origin(function(d) {
			return { x: gridScaleCol.rangeBand() / 2,
				     y: gridScaleRow.rangeBand() / 2 };
		})
    	.on('drag', dragMove)
		.on('dragstart', dragStart)
		.on('dragend', dragEnd);
    var dragColIndex = null;
    var dragRowIndex = null;
    
    // bound callbacks
    var onChange = [];
    var onRemove = [];
    var onAdd = [];
    
    // tooltip
    var cellTooltip = d3.tip()
        .attr('class', 'heatmap-tool-tip')
        .offset([-10, 0])
        .html(toolTipHtmlFunc)

    // dragging callback functions
    function dragStart(d) {
    	// silence other listeners while cell being dragged
    	d3.event.sourceEvent.stopPropagation();
    	// save initial row/col coordinates
    	dragColIndex = colMap.get(d.value.col).index;
    	dragRowIndex = rowMap.get(d.value.row).index;
    	// draw rectangles showing the current item being dragged
    	var g = rootSelection.select('svg').select('g');
    	g.append('rect').classed('drag-indicator col', true)
        	.attr('cursor', 'move')
    		.attr('x', gridScaleCol(dragColIndex))
    		.attr('y', 0)
    		.attr('width', gridScaleCol.rangeBand())
    		.attr('height', height);   		
    	g.append('rect').classed('drag-indicator row', true)
        	.attr('cursor', 'move')
    		.attr('x', 0)
    		.attr('y', gridScaleRow(dragRowIndex))
    		.attr('width', width)
    		.attr('height', gridScaleRow.rangeBand());
    }
    function dragMove(d) {
    	// compute distance dragged so far
    	var colDisp = Math.floor(d3.event.x / gridScaleCol.rangeBand());
    	var rowDisp = Math.floor(d3.event.y / gridScaleRow.rangeBand());
    	var colIndex = colMap.get(d.value.col).index;
    	var rowIndex = rowMap.get(d.value.row).index;    	
    	dragColIndex = Math.max(0, Math.min(colIds.length - 1, colIndex + colDisp));
    	dragRowIndex = Math.max(0, Math.min(rowIds.length - 1, rowIndex + rowDisp));
    	// update drag rectangle positions
    	var g = rootSelection.select('svg').select('g');
    	g.selectAll('.drag-indicator.col')
    	  .transition()
    	    .duration(50)
    		.attr('x', gridScaleCol(dragColIndex));
    	g.selectAll('.drag-indicator.row')
  	  	  .transition()
  	        .duration(50)
  		    .attr('y', gridScaleRow(dragRowIndex));
    }
    function dragEnd(d) {
    	// remove dragging rectangles
    	var g = rootSelection.select('svg').select('g');
    	g.selectAll('.drag-indicator').remove();
    	// issue move command
    	var colIndex = colMap.get(d.value.col).index;
    	var rowIndex = rowMap.get(d.value.row).index;    	
    	if (colIndex != dragColIndex) {
        	my.move(colIds, colMap, d.value.col, dragColIndex - colIndex);
    	}
    	if (rowIndex != dragRowIndex) {
        	my.move(rowIds, rowMap, d.value.row, dragRowIndex - rowIndex);
    	}
    	dragColIndex = null;
    	dragRowIndex = null;
    }
        
	// Private method to make row,col key for data map
	function make_key(r,c) {
	    return r + ',' + c;
	}
	function parse_key(k) {
		return k.split(',');
	}
    
    // Private method to reindex rol/column map after rows or columns 
    // are deleted
    function reindex(arr, amap) {
    	for (var i=0; i < arr.length; i++) {
    		amap.get(arr[i]).index = i;
    	}
    }
    
    // Private method to trigger event
    function trigger(funcArray, args) {
        for (var i=0; i < funcArray.length; i++) {
            funcArray[i](args);
        }    	
    }    
    
    // public properties and methods
    my = {};

	// event
    my.on = function(event, callback) {
		switch(event) {
		case 'add':
			onAdd.push(callback);
			break;
		case 'remove':
			onRemove.push(callback);
			break;
		case 'change':
			onChange.push(callback);			
			break;
		default:
			console.warn('unrecognized event' + event);
		}
	}

    // accessor for heatmap data values
    my.getData = function() {
    	return cellMap.entries();
    }
    
    // reset heatmap
    my.clear = function() {
    	delete rowIds;
    	rowIds = [];
    	delete colIds;
    	colIds = [];
    	delete rowMap;
        rowMap = d3.map();
    	delete colMap;
        colMap = d3.map();
    	delete cellMap;
        cellMap = d3.map();
        my.redraw();
    }

    // add rows/columns to heatmap
    my.add = function(studyName, rowEntries, colEntries) {
    	// rowEntries and colEntries are lists of objects where each object
    	// has a 'key' attribute and a 'value' object
    	var key, item;
    	var newRowIds = [];
    	var newColIds = [];    	
		for (var k=0; k < rowEntries.length; k++) {
			key = rowEntries[k].key;
			if (rowMap.has(key)) continue;
			newRowIds.push(key);
		}
        for (var k=0; k < colEntries.length; k++) {
			key = colEntries[k].key;
			if (colMap.has(key)) continue;
        	newColIds.push(key);
        }
        // request new data from server
        var d = JSON.stringify({ 
        	study: studyName,
            new_rows: newRowIds, 
            new_cols: newColIds,
            rows: rowIds,
            cols: colIds,
            value: heatmapValueField
        });
        d3.json(url)
        	.header("Content-Type", "application/json")
        	.post(d, function(error, resp) {
        		if (error) return console.warn(error);
        		// update rows and columns
        		for (var k=0; k < rowEntries.length; k++) {
                    // add new row
        			item = rowEntries[k];
        			if ( !rowMap.has(item.key) ) {
            			item.value.index = rowIds.length;
        				rowIds.push(item.key);
        				rowMap.set(item.key, item.value);
        			}
        		}
                for (var k=0; k < colEntries.length; k++) {
                    // add new column
                	item = colEntries[k];
                	if ( !colMap.has(item.key) ) {
                		item.value.index = colIds.length;
                		colIds.push(item.key);
                		colMap.set(item.key, item.value);
                	}
                }        		
        		// heatmap data
        		resp.data.forEach(function(el, index, array) {
        			cellMap.set(make_key(el.row, el.col), el);
        		});
        		my.redraw();
        		// trigger events
        		trigger(onAdd, { rows: newRowIds, cols: newColIds });
        		trigger(onChange);
        	});
    }

    // delete heatmap rows/cols    
	my.remove = function(removeRowIds, removeColIds) {
        // save old row and column id lists
        var oldRowIds = rowIds;
        var oldColIds = colIds;
    	// remove row data for all deleted rows
        for (var i=0; i < removeRowIds.length; i++) {
            rowMap.remove(removeRowIds[i]);
            // remove data for all columns in this row
            for (var j=0; j < oldColIds.length; j++) {
                cellMap.remove(make_key(removeRowIds[i], oldColIds[j]));
            }
        }   
        // remove column data for all deleted columns
        for (var j=0; j < removeColIds.length; j++) {
            colMap.remove(removeColIds[j]);
            // remove data for all rows in this column
            for (var i=0; i < oldRowIds.length; i++) {
            	cellMap.remove(make_key(oldRowIds[i], removeColIds[j]));
            }       
        }
        // remove row/col ids
        for (var i=0; i < removeRowIds.length; i++) {
        	var index = rowIds.indexOf(removeRowIds[i]);
        	rowIds.splice(index,1);
        }
        for (var j=0; j < removeColIds.length; j++) {
            var index = colIds.indexOf(removeColIds[j]);
            colIds.splice(index,1);
        }
        // reindex rowMap and colMap
        reindex(rowIds, rowMap);
        reindex(colIds, colMap);
        my.redraw();
        // trigger event
        trigger(onRemove, { rows: removeRowIds, cols: removeColIds });
        trigger(onChange);
	}
	
    // swap rows or columns
	my.move = function(arr, amap, id_from, amt) {
		// move element 'id' distance 'amt' within array and map
		// get index
		var index_from = amap.get(id_from).index;
		var index_to = index_from + amt;
		// check new index for valid move        
        if (index_to < 0) return;
        if (index_to >= arr.length) return;
        // get new id
        var id_to = arr[index_to];
        // swap indexes in rowMap or colMap
        amap.get(id_to).index = index_from;
        amap.get(id_from).index = index_to;
        // swap elements
		arr[index_to] = id_from;
		arr[index_from] = id_to;
	    // redraw
		my.redraw();
	}
	
	// sort array on row or column specified by axis and aid
	my.sort = function(axis, aid, reverse) {
        var sortMap;
        var sortIds;
        var keyFunc;
        if ( axis == 'row' ) {
			sortMap = colMap;
			sortIds = colIds;
			keyFunc = function(d) { return make_key(aid, d); }
        } else {
        	sortMap = rowMap;
        	sortIds = rowIds;
			keyFunc = function(d) { return make_key(d, aid); }
        }
        // loop over axis to sort on and get values
        var k, v;
        var sortData = [];
        for (var i=0; i < sortIds.length; i++) {
        	// lookup cell map
        	k = keyFunc(sortIds[i]);        	
        	v = cellMap.get(k).value;
        	sortData.push([sortIds[i],v]);
        }
        // sort decorated array
        sortData.sort(function(a, b) {
        	if ( reverse ) return (b[1] - a[1]);
        	return (a[1] - b[1]);
        });
    	// overwrite ids in new order and update map
        for (var i=0; i < sortIds.length; i++) {
        	k = sortData[i][0];
        	v = sortData[i][1];
        	sortIds[i] = k;
        	sortMap.get(k).index = i;
        }
        // redraw
        my.redraw();
	}
	
	my.cluster = function(axis) {
		var arr1;
		var amap;
		var arr2;
		var keyFunc;
		// cluster either rows or columns
        if ( axis == 'row' ) {
        	arr1 = rowIds;
        	amap = rowMap;
        	arr2 = colIds;
    		keyFunc = function(i,j) { return make_key(i,j); }
        } else {
        	arr1 = colIds;
        	amap = colMap;
        	arr2 = rowIds;
    		keyFunc = function(i,j) { return make_key(j,i); }
        }
        if ( !arr1.length ) return;
        if ( !arr2.length ) return;
		// build vector of values along axis
		var labels = [];
		var vectors = [];        
        var key, value;
        var vector;
		for (var i=0; i<arr1.length; i++) {
			vector = [];
			for (var j=0; j<arr2.length; j++) {
				key = keyFunc(arr1[i],arr2[j]);
				value = cellMap.get(key).value;
				vector.push(value);
			}
			labels[i] = arr1[i];
			vectors[i] = vector;
		}
		// perform clustering
		var root = figue.agglomerate(labels, vectors, figue.EUCLIDIAN_DISTANCE, figue.SINGLE_LINKAGE);
		var newlabels = []
		figue.flattenTree(root, newlabels);
    	// overwrite ids in new order and update map
        for (var i=0; i < newlabels.length; i++) {
        	arr1[i] = newlabels[i];
        	amap.get(newlabels[i]).index = i;
        }
	}
	
	// create the heatmap drawing using d3
	my.create = function(selection) {
        // save selection so we can easily redraw
        rootSelection = selection;
		// create main SVG
        var svg = selection.append("svg")
            .attr("width", width)
            .attr("height", height)
            .call(cellTooltip)
          .append('g')
            .classed('heatmap', true)
            .attr("transform", "translate(" + (padding + margin.left) + "," + (padding + margin.top) + ")");
        // bounding rectangle
        svg.append('rect')
            .style('stroke', 'gray')
            .style('fill', 'none')
            .attr('x', 0)
            .attr('y', 0)
            .attr('width', innerWidth)
            .attr('height', innerHeight);
        // clustering button
        svg.append('image')
        	.classed('cluster-button', true)
        	.attr('xlink:href', '/static/loop_32x32.png')
        	.attr('x', (-0.75)* margin.left)
        	.attr('y', (-0.75)* margin.top)
        	.attr('width', margin.left / 2)
        	.attr('height', margin.top / 2)
        	.on('click', function(d) {
	        	my.cluster('row');        		
	        	my.cluster('col');
	        	my.redraw();        		
        	});
	};
	
    my.redraw = function() {
        // redefine grid based on number of rows/cols in heatmap
        gridScaleRow.domain(d3.range(rowIds.length))
            .rangeBands([0, innerHeight]);
        gridScaleCol.domain(d3.range(colIds.length))
            .rangeBands([0, innerWidth]);
        
    	// select svg group that contains the plot
    	g = rootSelection.select('svg').select('g');  	
        
        // Row control buttons
        var rowButtons = g.selectAll('.row-button-group')
        	.data(rowIds, function(d) { return d; });
        var gRow = rowButtons.enter().append('g')
            .classed('row-button-group', true);
        gRow.append('image')
        	.classed('control-button', true)
        	.attr('xlink:href', '/static/x_alt_32x32.png')
        	.attr('x', -margin.left)
        	.attr('y', 0)
        	.attr("width", margin.left / numControlButtons)
        	.attr('cursor', 'pointer')
            .on('click', function(d) {
            	my.remove([d], []);
            });
        gRow.append('image')
    		.classed('control-button', true)
    		.attr('xlink:href', '/static/move_horizontal_alt2_32x32.png')
    		.attr('x', -margin.left + (margin.left / numControlButtons))
    		.attr('y', 0)
    		.attr("width", margin.left / numControlButtons)
    		.attr('cursor', 'pointer')
	        .on('click', function(d) {
	        	my.sort('row', d, toggle);
	        	toggle = !toggle;
	        });
        rowButtons
        	.transition().duration(duration)
        	.attr('transform', function(d, i) { 
				return 'translate(0,' + gridScaleRow(i) + ')';
			});
        rowButtons.selectAll('image')
        	.attr('height', gridScaleRow.rangeBand())
        rowButtons.exit().remove();        
        
        // Column controls (delete, move, sort)
        var colButtons = g.selectAll('.col-button-group')
    		.data(colIds, function(d) { return d; });
	    var gCol = colButtons.enter().append('g')
	        .classed('col-button-group', true);
	    gCol.append('image')
	    	.classed('control-button', true)
        	.attr('xlink:href', '/static/x_alt_32x32.png')
	    	.attr('x', 0)
	    	.attr('y', -margin.top)
	    	.attr('height', margin.top / numControlButtons)
	        .attr('cursor', 'pointer')
	        .on('click', function(d) {
	        	my.remove([], [d]);
	        });
	    gCol.append('image')
	    	.classed('control-button', true)
    		.attr('xlink:href', '/static/move_vertical_alt2_32x32.png')
	    	.attr('x', 0)
	    	.attr('y', -margin.top + (margin.top / numControlButtons))
	    	.attr('height', margin.top / numControlButtons)
	        .attr('cursor', 'pointer')
	        .on('click', function(d) {
	        	my.sort('col', d, toggle);
	        	toggle = !toggle;
	        });
	    colButtons
	    	.transition().duration(duration)
	    	.attr('transform', function(d, i) { 
				return 'translate(' + gridScaleCol(i) + ',0)';
			});
	    colButtons.selectAll('image')
	    	.attr('width', gridScaleCol.rangeBand())
	    colButtons.exit().remove();    

    	// Row Labels
        var rowLabels = g.selectAll(".rowLabel")
            .data(rowIds, function(d) { return d; });
        // Enter
        rowLabels.enter().append('text')
        	.classed('rowLabel', true)
            .style("text-anchor", "start")
            .attr("x", 6)
            .attr("dy", ".32em")
        	.each(function(d) {
        		return d3.select(this).call(rowLabelFunc, rowMap.get(d));
        	});
        // Enter + Update
        rowLabels.transition().duration(duration)
            .attr("y", gridScaleRow.rangeBand() / 2)
            .attr("transform", function(d, i) { return "translate(" + innerWidth + "," + gridScaleRow(i) + ")"; });
        // Exit
        rowLabels.exit().remove();
        
        // Column labels
        var colLabels = g.selectAll(".colLabel")
            .data(colIds, function(d) { return d; });
        // Enter
        colLabels.enter().append('text')
        	.classed('colLabel', true)
            .style("text-anchor", "end")
            .attr("x", -6)
            .attr("dy", ".32em");
        // Enter + Update
        colLabels.transition().duration(duration)
        	.text(function(d) {
                var v = colMap.get(d);
                if ( !v ) return "Loading..";
                return v[colLabelField];        
            })
            .attr("y", gridScaleCol.rangeBand() / 2)
            .attr("transform", function(d, i) { return "translate(" + gridScaleCol(i) + "," + innerHeight + ")rotate(-90)"; });    
        // Exit
        colLabels.exit().remove();
        
        // Cells
        var heatmap = g.selectAll(".cell")
            .data(cellMap.entries(), function(d) { return d.key; });
        // Enter
        heatmap.enter().append("rect")
        	.classed('cell', true)
        	.style('fill', 'white')
            .on('mouseover', cellTooltip.show)
            .on('mouseout', cellTooltip.hide)
        	.call(drag);
        // Enter + Update
        heatmap.transition().duration(duration)
            .attr("x", function(d) { return gridScaleCol(colMap.get(d.value.col).index); })
            .attr("y", function(d) { return gridScaleRow(rowMap.get(d.value.row).index); })
            .attr("width", gridScaleCol.rangeBand())
            .attr("height", gridScaleRow.rangeBand())
            .style("fill", function(d) {
                var v = d.value.value;
                if ( !v ) return colorDefault;
                return colorScale(v);           
            });
        // Exit
        heatmap.exit().remove();       
    }
    
    my.recolor = function() {
    	// select svg group that contains the plot
    	g = rootSelection.select('svg').select('g');
    	// recolor heatmap cells 
        g.selectAll(".cell")
            .data(cellMap.entries(), function(d) { return d.key; })
          .transition().duration(1000)
            .style("fill", function(d) {
                var v = d.value.value;
                if ( !v ) return colorDefault;
                return colorScale(v);           
            });
    }
	return my;
}

function HeatmapLegend(config) {
	// configuration
    var width = config.width || 300;
    var height = config.height || 300;
    var margin = config.margin || { top: 10, right: 10, bottom: 10, left: 10 };
    var axisOrient = config.axisOrient || 'top';
    var defaultColorDomain = config.colorDomain || [-20, 0, 20];
    var defaultColor = config.defaultColor || "#888888";
    var colorDomain = defaultColorDomain.slice(0);
    var colorRange = config.colorRange || ['blue', 'white', 'red'];

    // private variables
    var innerWidth = width - margin.left - margin.right;
    var innerHeight = height - margin.top - margin.bottom;
    var rootSelection = null;
    var onChange = [];
    
    // permanent state
    var colorScale = d3.scale.linear()
		.domain(colorDomain)
		.range(colorRange)
		.clamp(true);  
    var xScale = d3.scale.linear()
    	.domain([colorDomain[0],colorDomain[2]])
    	.range([0, innerWidth]);
    var xAxis = d3.svg.axis().scale(xScale).orient(axisOrient);    
    var yRandomizer = d3.random.normal(innerHeight / 2, innerHeight / 10); 
    
    // Private method to trigger event
    function trigger(funcArray) {
        for (var i=0; i < funcArray.length; i++) {
            funcArray[i]();
        }    	
    }
    
    // public methods and properties
    var my = {};

    // public colorscale accessor
    my.getColorScale = function() {
    	return colorScale;
    }
    
	// event
    my.on = function(event, callback) {
		switch(event) {
		case 'change':
			onChange.push(callback);			
			break;
		default:
			console.warn('unrecognized event' + event);
		}
	}
    
    // clear the legend drawing
    my.clear = function() {
    	// reset color scale and x axis scale
    	colorDomain = defaultColorDomain.slice(0);
    	colorScale.domain(colorDomain);
    	xScale.domain([colorDomain[0],colorDomain[2]]);
    }
    
	// create the legend drawing using d3
	my.create = function(selection) {
        // save selection so we can easily redraw
        rootSelection = selection;
		// Create the chart
        var svg = selection.append('svg')
            .attr("width", width)
            .attr("height", height);
        var g = svg.append('g')
        	.classed('heatmap-legend', true)
        	.attr("transform", "translate(" + margin.left + "," + margin.top + ")");
        // x axis
        g.append('g')
        	.classed('x axis', true)
        	.attr('transform', 'translate(0,' + innerHeight + ')')
        	.call(xAxis);        
        // selected domain rectangle
        g.append('rect')
        	.classed('select-box', true)
            .attr('pointer-events', 'none')
        // invisible rectangle to capture mouse clicks
        g.append('rect')
        	.classed('container', true)
            .attr("x", 0)
            .attr("y", 0)
            .attr("width", innerWidth)
            .attr("height", innerHeight)
            .attr('pointer-events', 'all')
            .on("click", function(d) {
                // handle mouse click in container
                var point = d3.mouse(this);
                var x = Math.abs(xScale.invert(point[0]));
                var domain = xScale.domain();
                left = Math.max(domain[0], -x);            
                right = Math.min(domain[1], x);
                colorDomain = [left, 0, right];
                colorScale.domain(colorDomain);
                redrawDomain();
                // notify event listeners
                trigger(onChange);
            });
        // draw domain
        redrawDomain();
	};
    
	// Update method requires data entries from heatmap
    my.update = function(data) {
    	// filter falsey values
    	data = data.filter(function (d) { return d.value.value; });
    	if ( !data.length ) return;
    	// find the extent of the data values
        var domain = d3.extent(data, function(d) { return d.value.value; });
        // update the x axis linear scale
        xScale.domain(domain);
        // Update the color scale to ensure it stays within the domain
        colorDomain[0] = Math.max(domain[0], colorDomain[0])
        colorDomain[2] = Math.min(domain[1], colorDomain[2])
        colorScale.domain(colorDomain);
    	// update the drawing
    	var g = rootSelection.select('svg').select('g');
    	// update the x axis
        g.select(".x.axis").call(xAxis);
        // update the data points
        var circle = g.selectAll('circle')
        	.data(data, function(d) { return d.key; });
        circle.enter().append('circle')
            .attr('pointer-events', 'none')
            .attr('r', 2)
            .attr('cy', yRandomizer());
        circle
            .attr("transform", function(d) { return "translate(" + xScale(d.value.value) + ",0)"; });
        circle.exit().remove();
        // Update the domain and colors
        redrawDomain();
	};
    
    function redrawDomain() {
        var g = rootSelection.select('svg').select('g');
        g.select('.select-box')
           .attr("x", xScale(colorDomain[0]))
           .attr("y", 0)
           .attr("width", xScale(colorDomain[2]) - xScale(colorDomain[0]))
           .attr("height", innerHeight);
        g.selectAll('circle')
    		.style('fill', function(d) {
    			var v = d.value.value;
    			if (!v) return defaultColor;
    			return colorScale(v);
    		});        
    };

    return my;    
}


